package HMM;

import java.util.ArrayList;

import TrainingData.TrainingInstance;
import TrainingData.TrainingInstanceCollection;


/**
 * @author Gergely Neu
 *
 * Policy matching training of an HMM.
 */
public class PolicyMatching extends TrainingMethod {

	protected double _beta = 0.0;
	
	/**
	 * Constructor for PolicyMatching. Extra parameter: beta.
	 * 
	 * @param iterNr number of iterations to carry out.
	 * @param trainingData the data used for training. 
	 * @param testData the data used for testing.
	 * @param beta the parameter of the Boltzmann distribution, the reciprocal of the temperature parameter.
	 */
	public PolicyMatching(int iterNr, TrainingInstanceCollection trainingData,TrainingInstanceCollection testData, double beta){
		super(iterNr,trainingData,testData);
		_beta = beta;
	}
	
	/**
	 * @author Gergely Neu
	 *
	 * A triple of integers used to count features in a decoded sequence.
	 *
	 */
	private class Triple {
		private int _prevState;
		private int _state;
		private int _emission;
		
		public Triple(int prevState, int state, int emission) {
			_prevState = prevState;
			_state = state;
			_emission = emission;
		}
		
		public int getPrevState() {
			return _prevState;
		}
		
		public int getState() {
			return _state;
		}
		
		public int getEmission() {
			return _emission;
		}
	}
	
	@Override
	protected Updates iterationStep(DoubleSparseMatrix transitionMx, DoubleSparseMatrix emissionMx) {
		
		int numStates = transitionMx.getColNr();
		int numEmissions = emissionMx.getColNr();
		
		DoubleSparseMatrix deltaTransMx = new DoubleSparseMatrix(numStates,numStates, 0.0);
		DoubleSparseMatrix deltaEmissMx = new DoubleSparseMatrix(numStates,numEmissions, 0.0);
		
		int inst = 0;
		
		for (TrainingInstance instance : _trainingData) { //TODO use counters! -- G.
			
			if ((inst++%100)==0 && false)
				System.out.println("[Training sentence "+ (inst-1) + "]");
			
			int[] emissions = instance.getEmissions();
			int[] states = instance.getStates();
			
			int L = emissions.length;
			
			HMMCommonFunctions.computeViterbiScores(emissions, emissionMx, transitionMx);
	
			for (int x = L-1; x>=1; x--) {
				//DoubleSparseMatrix pi = new DoubleSparseMatrix(numStates,1);
				double sumPi = 0.0;
				
				Counter<Triple> sumA = new Counter<Triple>();
				
				double constA = 0.0;
				
				double constB = 0.0;
				
				int action = 0;
				
				ArrayList<Double> pi = new ArrayList<Double>();
				
				for (int a : transitionMx.getPrecedingStates(states[x])) {
					double piA = Math.exp(_beta * HMMCommonFunctions.getViterbiScores(x-1, a));
					sumPi += piA;
					pi.add(action++, piA);
					
					int isCorrect = 0;
					
					if (a == states[x-1])
						isCorrect = 1;
					
					constA += isCorrect * piA;
					
					constB += piA * piA;
				}
				
				constB = constA / sumPi - constB / sumPi / sumPi;
				
				action = 0;
				
				for (int a : transitionMx.getPrecedingStates(states[x])) {

					double piA = pi.get(action++) / sumPi;
					
					if (piA == 0.0) continue;
					
					int isCorrect = 0;
					
					if (a == states[x-1])
						isCorrect = 1;
					
					int[] decodedStub = HMMCommonFunctions.decodeFrom(emissions, x, a);
					
					double update = (isCorrect - piA - constB) * piA; 
					
					for (int i = 0; i<x; i++) {
						int k = (i == 0) ? -1 : decodedStub[i-1];
						
						int l = decodedStub[i];
						
						int m = emissions[i];
						
						Triple newTriple = new Triple(k,l,m);
						
						sumA.incrementCount(newTriple, update);
					}
				}
				
				for (Triple triple : sumA.keySet()) {
					int k = triple.getPrevState();
					int l = triple.getState();
					int m = triple.getEmission();
					
					double count = sumA.getCount(triple);
					if (k>=0)
						deltaTransMx.increment(k, l, count);
					
					deltaEmissMx.increment(l, m, count);
				}
				
			}
		}
		
		return new Updates(deltaTransMx, deltaEmissMx);
	}
	
	
}
